import cv2
import numpy as np

from .utils import resize


class Mutator:
    """用于对图片做放射变换，以及对结果中的坐标做逆放射变换"""

    def __init__(self, target_size, crop_area=None, bi_directional=False):
        """
        :param target_size: Tuple[int, int]
            形如(height, width)的模型输入尺寸
        :param crop_area: Tuple[int, int, int, int]
            形如(left, top, right, bottom)的裁剪区域
        :param bi_directional: bool
            若需要坐标逆变换，则传入True，否则传入False

        Example:
            >>> t = Mutator(crop_area=(200, 300, 400, 500), target_size=(100, 150), bi_directional=True)
        """
        self.target_size = target_size
        self.crop_area = crop_area
        self.bi_directional = bi_directional and crop_area is not None
        if self.bi_directional:
            th, tw = target_size
            x1, y1, x2, y2 = crop_area
            src = np.array(
                [[x1, y1], [x2, y1], [x2, y2], [x1, y2]],
                dtype=np.float32,
            )
            dst = np.array([[0, 0], [tw, 0], [tw, th], [0, th]], dtype=np.float32)
            self.transposed_inv_mat = cv2.getPerspectiveTransform(dst, src).transpose()

    def mutate_image(self, image, for_test=False):
        """对图片做投影变换，随后转成模型输入形式: (c, h, w), float32

        Example:
            >>> img_org = np.zeros([600, 600, 3])
            >>> arr_trans = self.mutate_image(img_org)
            >>> print(arr_trans.shape, self.target_size)
            (3, 100, 150), (100, 150)
        """
        if self.crop_area is not None:
            x1, y1, x2, y2 = self.crop_area
            image = image[y1:y2, x1:x2]
        img_transformed = resize(image, (self.target_size[1], self.target_size[0]))
        input_arr = img_transformed.transpose(2, 0, 1).astype(np.float32)
        if for_test:
            input_arr = np.expand_dims(input_arr, 0)
        return input_arr

    def mutate_coordinates(self, coordinates):
        """对坐标做逆变换

        :param coordinates: np.ndarray
            形如[[X1, Y1], [x2, y2], ..., [xn, yn]]的ndarray坐标序列
        :return: np.ndarray
            与输入形式相同，顺序对应的逆变换后的坐标

        Example:
            >>> src = np.array([(0, 0), (0, 99), (149, 99), (149, 0)], dtype=np.float32)
            >>> print(self.mutate_coordinates(src).tolist())
            [[200, 300], [400, 300], [400, 500], [200, 500]]

        """
        src_expanded = np.concatenate([coordinates, np.ones([len(coordinates), 1])], 1)
        dst_expanded = src_expanded @ self.transposed_inv_mat
        return dst_expanded[:, :2]

    def trans_location(self, result):
        """对结果中的坐标进行逆变换，并返回整个结果"""
        if self.bi_directional:
            location = self.extract_location(result)
            location = self.mutate_coordinates(location)
            self.embed_location(result, location)
        return result

    def extract_location(self, result):
        """将坐标从结果中提取出来，形成n行2列的点列，每个点表示为(x,y)

        :param result:
            原始后处理结果
        :return:
            从原始后处理结果中提取出的坐标

        Example:
            >>> # 目标检测结果按照conf,bbox(xyxy),cls顺序排列时，按如下方法提取坐标
            >>> return result[:, 1:5].reshape(-1, 2)
        """
        if self.bi_directional:
            raise NotImplementedError("双向变换器必须实现此方法.")

    def embed_location(self, result, coordinates):
        """将变换后的坐标嵌入到结果中

        :param result:
            待嵌入结果
        :param coordinates:
            待嵌入坐标

        Example:
            >>> # 目标检测结果按照conf,bbox(xyxy),cls顺序排列时，按如下方法嵌入坐标
            >>> result[:, 1:5] = coordinates.reshape(-1, 4)
        """
        if self.bi_directional:
            raise NotImplementedError("双向变换器必须实现此方法")
